package oktx.dubbo.loadbalance;

import com.alibaba.dubbo.common.URL;
import com.alibaba.dubbo.common.extension.ExtensionLoader;
import com.alibaba.dubbo.rpc.Invocation;
import com.alibaba.dubbo.rpc.Invoker;
import com.alibaba.dubbo.rpc.cluster.LoadBalance;
import oktx.participant.JtaParticipant;
import oktx.utils.NetUtils;
import oktx.RemoteType;
import oktx.TransactionContext;

import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * Created by huangdachao on 2018/5/21 15:21.
 */
public class DubboTxLoadBalanceWrapper implements LoadBalance {
    private static volatile LoadBalance loadBalance;
    private LoadBalance instanceLoadBalance;

    public static void wrapStaticLoadBalance(LoadBalance lb) {
        loadBalance = lb;
    }

    public void wrapInstanceLoadBalance(LoadBalance lb) {
        this.instanceLoadBalance = lb;
    }

    private LoadBalance getLoadBalance() {
        if (instanceLoadBalance == null) {
            loadBalance = ExtensionLoader.getExtensionLoader(LoadBalance.class).getDefaultExtension();
            instanceLoadBalance = loadBalance;
        }
        return instanceLoadBalance;
    }

    @Override
    public <T> Invoker<T> select(List<Invoker<T>> invokers, URL url, Invocation invocation) {
        TransactionContext context = TransactionContext.getContext();
        if (context != null) {
            Set<Long> visited = context.getParticipants().stream().filter(p -> p instanceof JtaParticipant && p.getType() == RemoteType.dubbo)
                .map(p -> NetUtils.ipPortToNumber(p.getIp(), p.getPort())).collect(Collectors.toSet());

            Invoker<T> invoker = invokers.stream().filter(i -> visited.contains(NetUtils.ipPortToNumber(i.getUrl().getIp(), i.getUrl().getPort())))
                .findFirst().orElse(null);
            if (invoker != null) {
                return invoker;
            }
        }
        return getLoadBalance().select(invokers, url, invocation);
    }
}
